#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
完整的应急预案生成和保存示例
"""

import asyncio
import json
import aiohttp

async def generate_and_save_plan():
    """生成并保存应急预案的完整流程"""
    
    base_url = "http://localhost:8080"
    
    async with aiohttp.ClientSession() as session:
        print("=== 应急预案生成和保存完整流程 ===")
        print("=" * 50)
        
        # 步骤1：生成应急预案
        print("步骤1：生成应急预案")
        generate_url = f"{base_url}/api/v1/generate_plan"
        generate_data = {
            "event_type": "网络安全事件",
            "event_subtype": "数据勒索事件"
        }
        
        try:
            async with session.post(generate_url, json=generate_data) as response:
                if response.status == 200:
                    result = await response.json()
                    print(f"生成成功，状态码: {result.get('status_code')}")
                    
                    plan = result.get('plan', [])
                    if plan and isinstance(plan, list):
                        print("生成的应急预案内容：")
                        for item in plan:
                            print(f"【{item.get('field_name')}】")
                            print(f"{item.get('content')[:100]}...")
                        
                        # 步骤2：保存应急预案
                        print("\n步骤2：保存应急预案到知识库")
                        save_url = f"{base_url}/api/v1/save_plan"
                        save_data = {
                            "event_type": "网络安全事件",
                            "event_subtype": "数据勒索事件",
                            "plan": plan
                        }
                        
                        async with session.post(save_url, json=save_data) as save_response:
                            if save_response.status == 200:
                                save_result = await save_response.json()
                                print(f"保存成功，状态码: {save_result.get('status_code')}")
                                print(f"文件ID: {save_result.get('data', {}).get('file_id')}")
                                print(f"文件名: {save_result.get('data', {}).get('file_name')}")
                                print("应急预案已成功保存到知识库！")
                            else:
                                print(f"保存失败，状态码: {save_response.status}")
                                save_error = await save_response.text()
                                print(f"错误信息: {save_error}")
                    else:
                        print("生成的预案格式不正确")
                else:
                    print(f"生成失败，状态码: {response.status}")
                    error_text = await response.text()
                    print(f"错误信息: {error_text}")
                    
        except Exception as e:
            print(f"请求异常: {e}")

async def test_optimize_and_save():
    """测试优化并保存应急预案"""
    
    base_url = "http://localhost:8080"
    
    async with aiohttp.ClientSession() as session:
        print("\n=== 测试优化并保存应急预案 ===")
        print("=" * 50)
        
        # 步骤1：生成初始预案
        print("步骤1：生成初始预案")
        generate_url = f"{base_url}/api/v1/generate_plan"
        generate_data = {
            "event_type": "网络安全事件",
            "event_subtype": "数据勒索事件"
        }
        
        try:
            async with session.post(generate_url, json=generate_data) as response:
                if response.status == 200:
                    result = await response.json()
                    initial_plan = result.get('plan', [])
                    
                    if initial_plan:
                        # 步骤2：优化特定字段
                        print("步骤2：优化应急响应处置字段")
                        optimize_data = {
                            "event_type": "网络安全事件",
                            "event_subtype": "数据勒索事件",
                            "optimize_field": "应急响应处置",
                            "optimize_suggest": "增加具体的操作步骤和时间要求",
                            "history": [
                                {
                                    "role": "user",
                                    "content": "事件类型：网络安全事件，事件子类：数据勒索事件"
                                },
                                {
                                    "role": "assistant",
                                    "content": json.dumps({item["field_name"]: item["content"] for item in initial_plan}, ensure_ascii=False)
                                },
                                {
                                    "role": "user",
                                    "content": "优化'应急响应处置',优化建议：增加具体的操作步骤和时间要求"
                                }
                            ]
                        }
                        
                        async with session.post(generate_url, json=optimize_data) as optimize_response:
                            if optimize_response.status == 200:
                                optimize_result = await optimize_response.json()
                                optimized_content = optimize_result.get('plan', '')
                                
                                # 更新预案内容
                                updated_plan = []
                                for item in initial_plan:
                                    if item["field_name"] == "应急响应处置":
                                        updated_plan.append({
                                            "field_name": item["field_name"],
                                            "content": optimized_content
                                        })
                                    else:
                                        updated_plan.append(item)
                                
                                # 步骤3：保存优化后的预案
                                print("步骤3：保存优化后的预案")
                                save_data = {
                                    "event_type": "网络安全事件",
                                    "event_subtype": "数据勒索事件",
                                    "plan": updated_plan
                                }
                                
                                async with session.post(f"{base_url}/api/v1/save_plan", json=save_data) as save_response:
                                    if save_response.status == 200:
                                        save_result = await save_response.json()
                                        print(f"优化后的预案保存成功！")
                                        print(f"文件ID: {save_result.get('data', {}).get('file_id')}")
                                    else:
                                        print(f"保存失败，状态码: {save_response.status}")
                            else:
                                print(f"优化失败，状态码: {optimize_response.status}")
                    else:
                        print("初始预案生成失败")
                else:
                    print(f"初始预案生成失败，状态码: {response.status}")
                    
        except Exception as e:
            print(f"请求异常: {e}")

async def main():
    """主函数"""
    print("应急预案生成和保存完整示例")
    print("=" * 60)
    
    # 运行完整流程示例
    await generate_and_save_plan()
    
    # 运行优化并保存示例
    await test_optimize_and_save()

if __name__ == "__main__":
    asyncio.run(main()) 